#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time    : 2022/3/13 2:27 下午
# @Author  : WangZhixing


'''
使用方法：
python3 GraphDraw.py -r Datapath -o savepath
'''

import os
import sys

curPath = os.path.abspath(os.path.dirname(__file__))
rootPath = os.path.split(curPath)[0]
sys.path.append(rootPath)

import argparse
import warnings
from torch_geometric.transforms import NormalizeFeatures
from ProcessData import DependenceGraph

warnings.filterwarnings("ignore")
import matplotlib.pyplot as plt
import networkx as nx
from torch_geometric.utils.convert import to_networkx


def GraphDraw(Data, output=None):
    G = to_networkx(Data)
    H = G.to_undirected()
    pos = nx.kamada_kawai_layout(G)
    plt.figure(figsize=(19.20, 10.80))
    nx.draw(G,
            pos,
            with_labels=True,
            alpha=0.7,
            node_color=[[.7, .7, .7]],
            node_size=15,
            width=0.3,
            font_size=8,
            arrows=False
            )
    if output != None:
        plt.savefig(output)
    plt.show()
    plt.savefig(output)


def main():
    parser = argparse.ArgumentParser(description='this is root path')
    parser.add_argument('--root', '-r')
    parser.add_argument('--output', '-o', default=None)
    args = parser.parse_args()
    dataset = DependenceGraph(args.root, transform=NormalizeFeatures())
    data = dataset[0]
    GraphDraw(data, output=args.output)


if __name__ == "__main__":
    main()
